from tkinter import *
from tkinter.ttk import *
import tkinter
from ImageGridView import ImageGridView
import tkinter.filedialog as filedialog
from DataSet import DataSet
from LogUtils import logger
import os
import ImageUtils
from ClassSet import ClassSet


class SelectImageWindow(tkinter.Tk):
    def __init__(self):
        super().__init__()
        # self.attributes("-topmost", True)
        # self.dataSet = []
        # self.updateImageList(imageList)  # 主界面
        self.layoutWidgets()

    def updateImageList(self, imageList, trainNotSufList, testNotSufList):
        """
        更新图片列表
        :param imageList: 图片路径集合
        :return:
        """
        self.originImageList = []  # 未分配的完成图片了路径
        self.trainImageList = []  # 训练模块存储的图片完整路径
        self.testImageList = []  # 测试模块存储的图片完整路径
        logger.debug("传入的imageList{0}".format(imageList))
        for item in imageList:
            fileName = ImageUtils.getFileNameFromPath(item)  # 根据文件名来判别，没有其他方法
            if fileName in trainNotSufList:  # 如果图已经存在数据集中的训练集中
                self.trainImageList.append(item)
            elif fileName in testNotSufList:  # 如果图已经存在数据集中的测试集中
                self.testImageList.append(item)
            else:  # 剩下的就是未分配的图片
                self.originImageList.append(item)
        logger.debug("trainImageList{0}".format(self.trainImageList))
        logger.debug("testImageList{0}".format(self.testImageList))
        logger.debug("originImageList{0}".format(self.originImageList))
        self.originImageGridView.pushImagePathToImagePathFlags(self.originImageList)
        self.trainImageGridView.pushImagePathToImagePathFlags(self.trainImageList)
        self.testImageGridView.pushImagePathToImagePathFlags(self.testImageList)

    def layoutWidgets(self):
        self.title("挑选图片")
        s = tkinter.ttk.Style()
        s.configure('Treeview', rowheight=30)  # 设置TreeView的行高
        for x in range(4):
            for y in range(8):
                self.columnconfigure(y, weight=1, minsize=100)  # 设置列的最小宽度
                self.rowconfigure(x, weight=1, minsize=100)  # 设置行的最小宽度

        labelFrame = tkinter.LabelFrame(self, text="工具", width=100, labelanchor=N)
        labelFrame.grid(row=0, column=0, rowspan=4, padx=8, pady=8, sticky=W + S + N + E)
        for x in range(8):
            for y in range(1):
                labelFrame.columnconfigure(y, weight=1)  # 设置列的最小宽度
                labelFrame.rowconfigure(x, weight=1)  # 设置行的最小宽度
        topFrame = tkinter.Frame(labelFrame)
        topFrame.grid(row=0, column=0, rowspan=3, sticky=W + S + N + E)
        for x in range(4):
            for y in range(1):
                topFrame.columnconfigure(y, weight=1)  # 设置列的最小宽度
                topFrame.rowconfigure(x, weight=1)  # 设置行的最小宽度
        tkinter.Button(topFrame, font="Helvetic 10", text="打开图片集", command=self.openImageSet).grid(row=0, column=0,
                                                                                                   sticky=W + S + N + E,
                                                                                                   padx=8, pady=8)
        tkinter.Button(topFrame, font="Helvetic 10 bold", text="保存本次分类", command=self.saveClass).grid(row=1, column=0,
                                                                                                      sticky=W + S + N + E,
                                                                                                      padx=8, pady=8)
        tkinter.Button(topFrame, font="Helvetic 10 bold", text="生成VOC数据集", command=self.saveDataSet).grid(row=2,
                                                                                                          column=0,
                                                                                                          sticky=W + S + N + E,
                                                                                                          padx=8,
                                                                                                          pady=8)
        rangeLF = tkinter.LabelFrame(topFrame, text="随机选中")
        rangeLF.grid(row=3, column=0, sticky=W + N +E, padx=8, pady=8)
        self.progress = IntVar()
        self.progress.set(0)
        tkinter.Entry(rangeLF, textvariable=self.progress, bg="#34ace0", width=6).pack(side=TOP, fill=X)
        self.locationBtn = tkinter.Button(rangeLF, text="选中", width=6,command=self.randomSelect)
        self.locationBtn.pack(fill=X)

        bottomFrame = tkinter.Frame(labelFrame)
        bottomFrame.grid(row=3, column=0, rowspan=6, sticky=W + S + N + E)
        classListRT = Scrollbar(bottomFrame)
        classListRT.pack(side=RIGHT, fill=BOTH)
        self.classListTv = Treeview(bottomFrame, height=10, columns="classList", show="headings",
                                    yscrollcommand=classListRT.set)
        self.classListTv.pack(side=LEFT, fill=BOTH, expand=True)
        classListRT.config(command=self.classListTv.yview)
        self.classListTv.column("classList", anchor=CENTER, width=100)
        self.classListTv.heading("classList", text="分类文件")
        self.classListTv.tag_configure("imageFont", font="Helvetic 15 bold")
        self.classListTv.bind("<<TreeviewSelect>>", self.onSelectClass)
        self.originImageGridView = ImageGridView(self, title="未挑选的图片", rowCount=4, columnCount=2, length=150)
        self.originImageGridView.grid(row=0, column=1, rowspan=4, columnspan=4)

        self.trainImageGridView = ImageGridView(self, title="训练集", rowCount=2, columnCount=4, length=150, )
        self.trainImageGridView.grid(row=0, column=6, rowspan=2, columnspan=2)

        self.testImageGridView = ImageGridView(self, title="测试集", rowCount=2, columnCount=4, length=150)
        self.testImageGridView.grid(row=2, column=6, rowspan=2, columnspan=2)

        Button(self, text="------->", command=lambda: self.testImageGridView.pushImagePathToImagePathFlags(
            self.originImageGridView.popImagePathFromImagePathFlags())).grid(row=2, column=5, sticky=W + S + N + E,
                                                                             padx=8, pady=90)
        Button(self, text="<-------", command=lambda: self.originImageGridView.pushImagePathToImagePathFlags(
            self.testImageGridView.popImagePathFromImagePathFlags())).grid(row=3, column=5, sticky=W + S + N + E,
                                                                           padx=8, pady=90)
        Button(self, text="------->", command=lambda: self.trainImageGridView.pushImagePathToImagePathFlags(
            self.originImageGridView.popImagePathFromImagePathFlags())).grid(row=0, column=5, sticky=W + S + N + E,
                                                                             padx=8, pady=90)
        Button(self, text="<-------", command=lambda: self.originImageGridView.pushImagePathToImagePathFlags(
            self.trainImageGridView.popImagePathFromImagePathFlags())).grid(row=1, column=5, sticky=W + S + N + E,
                                                                            padx=8, pady=90)

    def randomSelect(self):
        rm = self.progress.get()
        self.originImageGridView.randomSelect(rm)

    def openImageSet(self):
        """
        打开已经分类的图片集
        :return:
        """
        path = filedialog.askdirectory()
        if os.path.exists(path):
            self.classificationImagePath = path  # 分类后的图片的根路径
            self.classList = ImageUtils.getDirFromDir(self.classificationImagePath)  # 读取路径下的文件 对等类型
            self.initData()
        pass

    def initData(self):
        """
        初始化数据
        :return:
        """
        for (index, item) in enumerate(self.classList):
            self.classListTv.insert("", iid=index, index=END, values=item,
                                    tags="imageFont")
        self.classListTv.selection_set(0)
        self.classListTv.see(0)

    def clearAllShow(self):
        self.originImageGridView.clearAllImage()
        self.trainImageGridView.clearAllImage()
        self.testImageGridView.clearAllImage()

    def onSelectClass(self, event):
        """
        响应事件
        :param event:
        :return:
        """
        logger.debug("是否响应事件")
        self.currentClassSet = None
        widgetObj = event.widget
        itemSelected = widgetObj.selection()[0]
        self.currentIndex = int(itemSelected)  # 选择的类型的index
        currentImagePath = self.classificationImagePath + "/" + self.classList[self.currentIndex]  # 当前选中的路径
        logger.debug("当前选中的路径为", currentImagePath)
        # self.loadCurrentClassImageSet(currentImagePath)

        self.currentClassSet = ClassSet(currentImagePath)
        self.clearAllShow()  # 清除显示
        self.originImageGridView.pushImagePathToImagePathFlags(self.currentClassSet.originImageList)
        self.trainImageGridView.pushImagePathToImagePathFlags(self.currentClassSet.trainImageList)
        self.testImageGridView.pushImagePathToImagePathFlags(self.currentClassSet.testImageList)

    def saveClass(self):
        """
        保存本次分类，比如保存preson的分类
        :return:
        """
        trainList = self.trainImageGridView.getAllImagePath()
        testList = self.testImageGridView.getAllImagePath()
        originList = self.originImageGridView.getAllImagePath()
        self.currentClassSet.changeImageFile(trainList,testList,originList)
        # self.currentClassSet.trainImageList = trainList
        # self.currentClassSet.testImageList = testList
        # self.currentClassSet.saveClass()
        pass

    def initEvent(self):
        """
        初始化点击事件
        :return:
        """

        pass

    def saveDataSet(self):
        savePath = filedialog.askdirectory()

        # TODO:需要改变
        dataSet = DataSet(savePath)
        trainSet = []
        testSet = []
        for item in self.classList:
            trainPath = self.classificationImagePath + "/" + item + "/train"
            testPath = self.classificationImagePath + "/" + item + "/test"
            # print(ImageUtils.getSpecifiedTypeFileFromDirectory(trainPath, [".jpg", ".png"]))
            # print(ImageUtils.getSpecifiedTypeFileFromDirectory(testPath, [".jpg", ".png"]))
            trainSet = trainSet + ImageUtils.getSpecifiedTypeFileFromDirectory(trainPath, [".jpg", ".png",".jpeg"])
            testSet = testSet + ImageUtils.getSpecifiedTypeFileFromDirectory(testPath, [".jpg", ".png",".jpeg"])
        logger.debug("打印测试集", testSet,"共有",len(testSet))
        logger.debug("打印训练集", trainSet,"共有",len(trainSet))
        for item in trainSet:
            dataSet.pushTrainSetImage(item)
        for item in testSet:
            dataSet.pushTestSetImage(item)
        # dataSet.writeClassName(self.classList)
        dataSet.saveDataSet(self.classificationImagePath)
        pass

    # def loadCurrentClassImageSet(self, classPath):
    #     """
    #     加载当前类型的图片集
    #     :param classPath: 类型所在路径
    #     :return:
    #     """
    #     trainPath = classPath + "/train"
    #     testPath = classPath + "/test"
    #     ImageUtils.check_and_create_dir(trainPath)
    #     ImageUtils.check_and_create_dir(testPath)
    #     self.originImageGridView.clearAllImage()
    #     self.testImageGridView.clearAllImage()
    #     self.trainImageGridView.clearAllImage()
    #     originImageNameList = ImageUtils.getFilePathFromDir(classPath,[".jpg",".png"])  # 原始图片，完整路径
    #     trainImageNameList = ImageUtils.getFileNameNotSufFromDir(trainPath, [".xml", ".png"])  # 训练图片，只包含文件名
    #     testImageNameList = ImageUtils.getFileNameNotSufFromDir(testPath, [".xml", ".png"])  # 测试图片，只包含文件名
    #     self.updateImageList(originImageNameList, trainImageNameList, testImageNameList)
    #     pass


if __name__ == '__main__':
    window = SelectImageWindow()
    # window.attributes('-fullscreen', True)
    w = window.winfo_screenwidth()
    h = window.winfo_screenheight()
    window.geometry("%dx%d" % (w, h))
    window.mainloop()
    pass
